import torch.nn as nn

# LSTM模型
class LSTMModel(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.args = args
        self.lstm = nn.LSTM(
            input_size=args.input_size,
            hidden_size=args.hidden_size,
            num_layers=args.num_layers,
            batch_first=True
        )
        self.linear = nn.Linear(args.hidden_size, args.output_size)
        
    def forward(self, x):
        out, _ = self.lstm(x)
        out = self.linear(out[:, -1, :])
        return out